import tkinter as tk
from tkinter import ttk, messagebox
import paramiko
import threading
import time
import socket
import sys
import json
import os


class SSHTunnelManager:
    def __init__(self, root):
        self.root = root
        self.root.title("SSH Tunnel Manager")
        self.root.geometry("600x700")

        # SSH连接信息
        self.ssh_config = {
            'hostname': tk.StringVar(),
            'username': tk.StringVar(),
            'password': tk.StringVar(),
            'port': tk.StringVar(value="22")  # 添加SSH端口，默认22
        }

        # 端口映射列表
        self.tunnel_mappings = []

        # 创建UI界面
        self.create_widgets()

        # 加载配置文件
        self.load_config()

    def create_widgets(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

        # 配置SSH信息区域
        ssh_frame = ttk.LabelFrame(main_frame, text="SSH连接配置", padding="10")
        ssh_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=5)

        # 主机IP输入
        ttk.Label(ssh_frame, text="主机IP:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=5)
        ttk.Entry(ssh_frame, textvariable=self.ssh_config['hostname'], width=30).grid(row=0, column=1,
                                                                                      sticky=(tk.W, tk.E), padx=5,
                                                                                      pady=5)

        # SSH端口输入
        ttk.Label(ssh_frame, text="SSH端口:").grid(row=1, column=0, sticky=tk.W, padx=5, pady=5)
        ttk.Entry(ssh_frame, textvariable=self.ssh_config['port'], width=30).grid(row=1, column=1,
                                                                                  sticky=(tk.W, tk.E), padx=5,
                                                                                  pady=5)

        # 用户名输入
        ttk.Label(ssh_frame, text="用户名:").grid(row=2, column=0, sticky=tk.W, padx=5, pady=5)
        ttk.Entry(ssh_frame, textvariable=self.ssh_config['username'], width=30).grid(row=2, column=1,
                                                                                      sticky=(tk.W, tk.E), padx=5,
                                                                                      pady=5)

        # 密码输入
        ttk.Label(ssh_frame, text="密码:").grid(row=3, column=0, sticky=tk.W, padx=5, pady=5)
        ttk.Entry(ssh_frame, textvariable=self.ssh_config['password'], show="*", width=30).grid(row=3, column=1,
                                                                                                sticky=(tk.W, tk.E),
                                                                                                padx=5, pady=5)

        # 端口映射区域
        mapping_frame = ttk.LabelFrame(main_frame, text="端口映射配置", padding="10")
        mapping_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=5)

        # 本地端口输入
        ttk.Label(mapping_frame, text="本地端口:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=5)
        self.local_port_entry = ttk.Entry(mapping_frame, width=15)
        self.local_port_entry.grid(row=0, column=1, sticky=tk.W, padx=5, pady=5)

        # 远程端口输入
        ttk.Label(mapping_frame, text="远程端口:").grid(row=0, column=2, sticky=tk.W, padx=5, pady=5)
        self.remote_port_entry = ttk.Entry(mapping_frame, width=15)
        self.remote_port_entry.grid(row=0, column=3, sticky=tk.W, padx=5, pady=5)

        # 添加按钮
        add_btn = ttk.Button(mapping_frame, text="添加映射", command=self.add_mapping)
        add_btn.grid(row=0, column=4, padx=5, pady=5)

        # 映射列表显示区域
        list_frame = ttk.LabelFrame(main_frame, text="端口映射列表", padding="10")
        list_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), pady=5)

        # 创建Treeview显示映射列表
        columns = ('本地端口', '远程端口')
        self.mapping_listbox = ttk.Treeview(list_frame, columns=columns, show='headings', height=8)
        self.mapping_listbox.heading('本地端口', text='本地端口')
        self.mapping_listbox.heading('远程端口', text='远程端口')

        # 添加滚动条
        scrollbar = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.mapping_listbox.yview)
        self.mapping_listbox.configure(yscrollcommand=scrollbar.set)

        self.mapping_listbox.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))

        # 删除按钮
        delete_btn = ttk.Button(list_frame, text="删除选中", command=self.delete_mapping)
        delete_btn.grid(row=1, column=0, pady=5)

        # 清空按钮
        clear_btn = ttk.Button(list_frame, text="清空列表", command=self.clear_mappings)
        clear_btn.grid(row=1, column=1, pady=5)

        # 连接控制区域
        control_frame = ttk.LabelFrame(main_frame, text="连接控制", padding="10")
        control_frame.grid(row=3, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=5)

        # 连接按钮
        self.connect_btn = ttk.Button(control_frame, text="建立连接", command=self.start_tunnel)
        self.connect_btn.grid(row=0, column=0, padx=5, pady=5)

        # 断开按钮
        self.disconnect_btn = ttk.Button(control_frame, text="断开连接", command=self.stop_tunnel, state=tk.DISABLED)
        self.disconnect_btn.grid(row=0, column=1, padx=5, pady=5)

        # 状态显示区域
        self.status_label = ttk.Label(control_frame, text="状态: 未连接")
        self.status_label.grid(row=1, column=0, columnspan=2, pady=5)

        # 配置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(0, weight=1)
        main_frame.rowconfigure(2, weight=1)
        mapping_frame.columnconfigure(1, weight=1)
        list_frame.columnconfigure(0, weight=1)
        list_frame.rowconfigure(0, weight=1)

        # SSH客户端实例
        self.ssh_client = None
        self.tunnel_threads = []
        self.tunnel_channels = []

    def add_mapping(self):
        local_port = self.local_port_entry.get()
        remote_port = self.remote_port_entry.get()

        if not local_port or not remote_port:
            messagebox.showwarning("警告", "请填写本地端口和远程端口")
            return

        try:
            local_port_int = int(local_port)
            remote_port_int = int(remote_port)

            # 检查端口是否重复
            for item in self.tunnel_mappings:
                if item['local'] == local_port_int and item['remote'] == remote_port_int:
                    messagebox.showwarning("警告", "该映射已存在")
                    return

            mapping = {
                'local': local_port_int,
                'remote': remote_port_int
            }

            self.tunnel_mappings.append(mapping)
            self.mapping_listbox.insert('', tk.END, values=(local_port, remote_port))

            # 清空输入框
            self.local_port_entry.delete(0, tk.END)
            self.remote_port_entry.delete(0, tk.END)

        except ValueError:
            messagebox.showerror("错误", "端口号必须是数字")

    def delete_mapping(self):
        selected_items = self.mapping_listbox.selection()
        if not selected_items:
            messagebox.showwarning("警告", "请选择要删除的映射项")
            return

        for item in selected_items:
            values = self.mapping_listbox.item(item, 'values')
            local_port = int(values[0])
            remote_port = int(values[1])

            # 从列表中删除
            for mapping in self.tunnel_mappings:
                if mapping['local'] == local_port and mapping['remote'] == remote_port:
                    self.tunnel_mappings.remove(mapping)
                    break

            # 从Treeview中删除
            self.mapping_listbox.delete(item)

    def clear_mappings(self):
        self.tunnel_mappings.clear()
        for item in self.mapping_listbox.get_children():
            self.mapping_listbox.delete(item)

    def start_tunnel(self):
        hostname = self.ssh_config['hostname'].get()
        username = self.ssh_config['username'].get()
        password = self.ssh_config['password'].get()
        port = self.ssh_config['port'].get()

        if not hostname or not username or not password:
            messagebox.showwarning("警告", "请填写完整的SSH连接信息")
            return

        if not self.tunnel_mappings:
            messagebox.showwarning("警告", "请至少添加一个端口映射")
            return

        # 验证SSH端口
        try:
            port_int = int(port)
            if port_int < 1 or port_int > 65535:
                messagebox.showerror("错误", "SSH端口号必须在1-65535之间")
                return
        except ValueError:
            messagebox.showerror("错误", "SSH端口号必须是数字")
            return

        # 禁用连接按钮，启用断开按钮
        self.connect_btn.config(state=tk.DISABLED)
        self.disconnect_btn.config(state=tk.NORMAL)
        self.status_label.config(text="状态: 建立连接中...")

        # 在新线程中启动隧道
        thread = threading.Thread(target=self._create_tunnels, args=(hostname, username, password, port))
        thread.daemon = True
        thread.start()

    def _create_tunnels(self, hostname, username, password, port):
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.ssh_client.connect(hostname, username=username, password=password, port=int(port))

            # 为每个映射创建隧道
            for mapping in self.tunnel_mappings:
                local_port = mapping['local']
                remote_port = mapping['remote']

                thread = threading.Thread(target=self._create_single_tunnel, args=(local_port, remote_port))
                thread.daemon = True
                thread.start()
                self.tunnel_threads.append(thread)

            self.root.after(0, lambda: self.status_label.config(text="状态: 已连接"))

        except Exception as e:
            self.root.after(0, lambda: messagebox.showerror("错误", f"SSH连接失败: {str(e)}"))
            self.root.after(0, lambda: self.connect_btn.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.disconnect_btn.config(state=tk.DISABLED))

    def _create_single_tunnel(self, local_port, remote_port):
        try:
            # 创建SSH隧道
            transport = self.ssh_client.get_transport()

            # Create a socket for the local port
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.bind(('localhost', local_port))
            sock.listen(1)

            # Store channel and socket for cleanup
            tunnel_info = {
                'local_port': local_port,
                'remote_port': remote_port,
                'socket': sock,
                'channel': None
            }

            self.tunnel_channels.append(tunnel_info)

            while True:
                # Accept incoming connections on the local port
                client_socket, addr = sock.accept()

                # Create tunnel channel for each connection
                channel = transport.open_channel("direct-tcpip", ("localhost", remote_port), ("127.0.0.1", local_port))

                # Store the channel reference
                tunnel_info['channel'] = channel

                # Handle data transfer in a separate thread
                threading.Thread(target=self._handle_tunnel_data, args=(client_socket, channel)).start()

        except Exception as e:
            self.root.after(0, lambda: messagebox.showerror("错误", f"隧道创建失败: {str(e)}"))
            print(f"隧道创建失败: {str(e)}")

    def _handle_tunnel_data(self, client_socket, channel):
        """Handle data transfer between local and remote sockets"""
        try:
            # Forward data between the sockets
            while True:
                # Read from client socket
                data = client_socket.recv(1024)
                if not data:
                    break

                # Send to SSH channel (remote side)
                channel.send(data)

                # Read from SSH channel and send back to client
                data = channel.recv(1024)
                if not data:
                    break

                client_socket.send(data)

        except Exception as e:
            print(f"数据转发错误: {str(e)}")
        finally:
            try:
                client_socket.close()
                channel.close()
            except:
                pass

    def stop_tunnel(self):
        # Close all tunnel channels
        for tunnel_info in self.tunnel_channels:
            try:
                if tunnel_info['channel']:
                    tunnel_info['channel'].close()
                tunnel_info['socket'].close()
            except Exception as e:
                print(f"关闭隧道时出错: {str(e)}")

        # Clear stored channels
        self.tunnel_channels.clear()

        # Stop all threads
        for thread in self.tunnel_threads:
            try:
                thread.join(timeout=1)
            except:
                pass

        self.tunnel_threads.clear()

        if self.ssh_client:
            self.ssh_client.close()

        # 启用连接按钮，禁用断开按钮
        self.connect_btn.config(state=tk.NORMAL)
        self.disconnect_btn.config(state=tk.DISABLED)
        self.status_label.config(text="状态: 已断开")

    def on_closing(self):
        self.stop_tunnel()
        self.save_config()  # Save configuration before closing
        self.root.destroy()

    def save_config(self):
        config_data = {
            'ssh': {
                'hostname': self.ssh_config['hostname'].get(),
                'username': self.ssh_config['username'].get(),
                'password': self.ssh_config['password'].get(),  # Note: Storing passwords in plain text is not secure
                'port': self.ssh_config['port'].get()
            },
            'mappings': self.tunnel_mappings
        }

        with open('config.json', 'w') as config_file:
            json.dump(config_data, config_file, indent=4)

    def load_config(self):
        if os.path.exists('config.json'):
            try:
                with open('config.json', 'r') as config_file:
                    config_data = json.load(config_file)
                    self.ssh_config['hostname'].set(config_data['ssh']['hostname'])
                    self.ssh_config['username'].set(config_data['ssh']['username'])
                    self.ssh_config['password'].set(config_data['ssh']['password'])  # Note: Storing passwords in plain text is not secure
                    self.ssh_config['port'].set(config_data['ssh']['port'])

                    for mapping in config_data['mappings']:
                        self.tunnel_mappings.append(mapping)
                        self.mapping_listbox.insert('', tk.END, values=(mapping['local'], mapping['remote']))

            except Exception as e:
                messagebox.showerror("错误", f"加载配置文件失败: {str(e)}")
        else:
            # Create an empty config.json file if it doesn't exist
            with open('config.json', 'w') as config_file:
                json.dump({}, config_file, indent=4)


if __name__ == "__main__":
    root = tk.Tk()
    app = SSHTunnelManager(root)

    # 设置窗口关闭事件
    root.protocol("WM_DELETE_WINDOW", app.on_closing)

    root.mainloop()